package com.huawei.spare.part.replace.utils;

import lombok.Getter;
import net.schmizz.keepalive.KeepAliveProvider;
import net.schmizz.sshj.DefaultConfig;
import net.schmizz.sshj.SSHClient;
import net.schmizz.sshj.common.LoggerFactory;
import net.schmizz.sshj.common.StreamCopier;
import net.schmizz.sshj.connection.channel.direct.Session;
import net.schmizz.sshj.transport.verification.PromiscuousVerifier;
import org.apache.commons.io.IOUtils;
import org.slf4j.Logger;

import java.io.ByteArrayOutputStream;
import java.io.Closeable;
import java.io.IOException;
import java.util.concurrent.TimeUnit;

@Getter
public class PTY implements Closeable {

    private static final Logger log = org.slf4j.LoggerFactory.getLogger(PTY.class);

    public static final String CHARSET = "UTF-8";

    SSHClient sshClient;
    Session session;
    Session.Shell shell;
    ByteArrayOutputStream bos;

    private DefaultConfig getKeepAliveConfig() {
        DefaultConfig config = new DefaultConfig();
        config.setKeepAliveProvider(KeepAliveProvider.KEEP_ALIVE);
        return config;
    }

    public static PTY allocate(String hostname, String username, String password) throws IOException {
        PTY $this = new PTY();
        $this.sshClient = new SSHClient($this.getKeepAliveConfig());
        $this.sshClient.addHostKeyVerifier(new PromiscuousVerifier());
        $this.sshClient.connect(hostname);
        $this.sshClient.authPassword(username, password);
        $this.sshClient.getConnection().getKeepAlive().setKeepAliveInterval(3);

        $this.session = $this.sshClient.startSession();
        $this.session.allocateDefaultPTY();
        $this.shell = $this.session.startShell();

        $this.listen();
        return $this;
    }

    public void listen() {
        bos = new ByteArrayOutputStream();
        new StreamCopier(shell.getInputStream(), bos, LoggerFactory.DEFAULT)
                .bufSize(shell.getRemoteMaxPacketSize())
                .keepFlushing(true)
                .spawn("PTY-read-thread");
    }

    public void send(String content) throws IOException {
        shell.getOutputStream().write(content.getBytes(CHARSET));
        shell.getOutputStream().flush();
    }

    @Override
    public void close() throws IOException {
        shell.close();
        session.close();
        sshClient.disconnect();
        sshClient.close();
    }

    public static void main(String[] args) throws IOException {
        try (PTY pty = PTY.allocate("115.159.160.190", "root", "!QAZ2wsx#EDC")) {
            pty.send("python /root/spare.py\n");
//            String s = IOUtils.toString(pty.bos.toByteArray(), CHARSET);
            System.out.println( pty.bos.toString());

            pty.send("1\n");
            String s2 = IOUtils.toString(pty.bos.toByteArray(), CHARSET);
            System.out.println(s2);
            pty.send("2\n");
//            String s2 = IOUtils.toString(pty.bos.toByteArray(), CHARSET);
            System.out.println( pty.bos.toString());
            pty.send("3\n");
//            String s2 = IOUtils.toString(pty.bos.toByteArray(), CHARSET);
            System.out.println( pty.bos.toString());
//        String read2 = pty.read();
//        pty.send("Choice 2");
//        String read3 = pty.read();
//        pty.send("Choice 3");
//        String read4 = pty.read();

            System.out.println();
        }

    }
}
